import yaml
import json
import logging
from datetime import datetime
from pathlib import Path
from pymongo import MongoClient
from docx import Document
from docx.shared import Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
import re
import requests
from urllib.parse import urlparse
import os
from PIL import Image
import io

class DocxGenerator:
    def __init__(self):
        self.config = self.load_config()
        self.setup_logging()
        self.setup_mongodb()
        self.clear_temp_images()
        
    def setup_logging(self):
        """设置日志"""
        log_dir = Path(__file__).parent / 'logs'
        log_dir.mkdir(exist_ok=True)
        
        log_filename = log_dir / f"docx_generator_{datetime.now().strftime('%Y%m%d')}.log"
        
        self.logger = logging.getLogger('DocxGenerator')
        self.logger.setLevel(logging.INFO)
        
        if not self.logger.handlers:
            file_handler = logging.FileHandler(log_filename, encoding='utf-8')
            console_handler = logging.StreamHandler()
            
            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
            file_handler.setFormatter(formatter)
            console_handler.setFormatter(formatter)
            
            self.logger.addHandler(file_handler)
            self.logger.addHandler(console_handler)

    def setup_mongodb(self):
        """设置MongoDB连接"""
        try:
            db_config = self.config.get('database', {})
            host = db_config.get('host', 'localhost')
            port = db_config.get('port', 27017)
            db_name = db_config.get('db_name', 'InfoStore')
            
            self.client = MongoClient(host, port)
            self.db = self.client[db_name]
            self.logger.info(f"MongoDB连接成功: {host}:{port}/{db_name}")
        except Exception as e:
            self.logger.error(f"MongoDB连接失败: {str(e)}")
            self.client = None
            self.db = None

    def clear_temp_images(self):
        """清空临时图片文件夹"""
        try:
            temp_dir = Path(__file__).parent / 'temp_images'
            if temp_dir.exists():
                # 删除文件夹中的所有文件
                for file_path in temp_dir.iterdir():
                    if file_path.is_file():
                        try:
                            file_path.unlink()
                            self.logger.info(f"已删除临时图片文件: {file_path}")
                        except Exception as e:
                            self.logger.warning(f"删除文件失败: {file_path}, 错误: {str(e)}")
                
                self.logger.info("临时图片文件夹已清空")
            else:
                self.logger.info("临时图片文件夹不存在，无需清空")
        except Exception as e:
            self.logger.error(f"清空临时图片文件夹失败: {str(e)}")

    def load_config(self):
        """加载配置文件"""
        config_path = Path(__file__).parent / 'config.yaml'
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            self.logger.error(f"配置文件未找到: {config_path}")
            return {}
        except Exception as e:
            self.logger.error(f"加载配置文件失败: {str(e)}")
            return {}

    def generate_docx_files(self):
        """生成docx文件的主函数"""
        if self.db is None:
            self.logger.error("MongoDB连接失败，无法生成文档")
            return
            
        try:
            # 获取所有未处理的Article_Info记录
            article_info_collection = self.db['Article_Info']
            unprocessed_articles = article_info_collection.find({
                'processed': {'$ne': True}  # 查找未处理的文章
            })
            
            processed_count = 0
            for article_info in unprocessed_articles:
                try:
                    self.generate_single_docx(article_info)
                    processed_count += 1
                except Exception as e:
                    self.logger.error(f"生成文档失败，文章ID: {article_info.get('dd_article_id')}, 错误: {str(e)}")
            
            self.logger.info(f"成功生成 {processed_count} 个docx文件")
            
        except Exception as e:
            self.logger.error(f"生成docx文件过程中发生错误: {str(e)}")

    def extract_text(self, data, result=None):
        """提取文本内容
        处理规则：
        1. 如果元素是字典且包含title键，输出"【标题】" + title值
        2. 如果元素是字典且type不是audio和salutation，输出其text内容
        3. 如果type是image，处理图片下载
        """
        if result is None:
            result = []
        
        if not isinstance(data, list):
            return result
            
        for item in data:
            if not isinstance(item, dict):
                continue
                
            # 处理标题
            #if 'title' in item:
            #    result.append(f"【标题】{item['title']}")
                
            # 处理文本内容
            if 'type' in item and item['type'] not in ['audio', 'salutation']:
                if item['type'] == 'image' and 'url' in item:
                    # 处理图片
                    image_path = self.download_image(item['url'])
                    if image_path:
                        result.append({'type': 'image', 'path': image_path})
                elif 'text' in item:
                    result.append(item['text'])
        return result

    def download_image(self, url):
        """下载图片并返回本地路径"""
        try:
            # 解析URL获取文件名和扩展名
            parsed_url = urlparse(url)
            original_filename = os.path.basename(parsed_url.path)
            file_ext = os.path.splitext(original_filename)[1].lower()
            
            # 设置请求头，模拟浏览器访问
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Accept': 'image/webp,image/apng,image/*,*/*;q=0.8',
                'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
                'Accept-Encoding': 'gzip, deflate, br',
                'Connection': 'keep-alive',
                'Upgrade-Insecure-Requests': '1'
            }
            
            # 生成唯一文件名
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')[:-3]  # 精确到毫秒
            unique_id = hash(url) % 10000  # 使用URL哈希生成短ID
            
            # 确定文件扩展名
            if file_ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff']:
                filename = f"img_{timestamp}_{unique_id}{file_ext}"
            else:
                # 如果没有扩展名，尝试从Content-Type获取
                try:
                    head_response = requests.head(url, headers=headers, timeout=10, verify=False)
                    content_type = head_response.headers.get('content-type', '')
                    if 'image/jpeg' in content_type:
                        filename = f"img_{timestamp}_{unique_id}.jpg"
                    elif 'image/png' in content_type:
                        filename = f"img_{timestamp}_{unique_id}.png"
                    elif 'image/gif' in content_type:
                        filename = f"img_{timestamp}_{unique_id}.gif"
                    else:
                        filename = f"img_{timestamp}_{unique_id}.jpg"
                except Exception as e:
                    self.logger.warning(f"获取图片信息失败: {str(e)}, 使用默认文件名")
                    filename = f"img_{timestamp}_{unique_id}.jpg"
            
            # 创建临时目录
            temp_dir = Path(__file__).parent / 'temp_images'
            temp_dir.mkdir(exist_ok=True)
            
            # 下载图片
            self.logger.info(f"开始下载图片: {url}")
            response = requests.get(url, headers=headers, timeout=30, verify=False, stream=True)
            response.raise_for_status()
            
            # 获取Content-Length
            content_length = response.headers.get('content-length')
            if content_length:
                self.logger.info(f"服务器报告文件大小: {content_length} 字节")
            
            # 保存图片
            image_path = temp_dir / filename
            downloaded_size = 0
            with open(image_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        downloaded_size += len(chunk)
            
            # 验证文件是否成功保存
            if os.path.exists(image_path) and os.path.getsize(image_path) > 0:
                actual_size = os.path.getsize(image_path)
                self.logger.info(f"下载完成 - 实际大小: {actual_size} 字节, 下载大小: {downloaded_size} 字节")
                
                # 如果服务器报告了大小，检查是否匹配
                if content_length and actual_size != int(content_length):
                    self.logger.warning(f"文件大小不匹配 - 服务器: {content_length}, 实际: {actual_size}")
                
                # 验证图片格式
                try:
                    with Image.open(image_path) as img:
                        self.logger.info(f"图片验证通过 - 格式: {img.format}, 尺寸: {img.size}")
                        return str(image_path)
                except Exception as e:
                    self.logger.error(f"下载的图片格式验证失败: {str(e)}")
                    # 尝试重新下载一次
                    self.logger.info("尝试重新下载...")
                    return self.download_image(url)
            else:
                self.logger.error(f"图片文件保存失败或文件为空: {image_path}")
                return None
            
        except requests.exceptions.RequestException as e:
            self.logger.error(f"图片下载网络错误: {url}, 错误: {str(e)}")
            return None
        except Exception as e:
            self.logger.error(f"图片下载失败: {url}, 错误: {str(e)}")
            return None

    def validate_and_fix_image(self, image_path):
        """验证并修复图片格式，确保与python-docx兼容"""
        try:
            self.logger.info(f"验证图片格式: {image_path}")
            
            # 尝试用PIL打开图片
            with Image.open(image_path) as img:
                # 检查图片格式
                format_name = img.format
                mode = img.mode
                size = img.size
                
                self.logger.info(f"图片信息 - 格式: {format_name}, 模式: {mode}, 尺寸: {size}")
                
                # 强制重新保存图片以确保兼容性
                # 生成新的文件名
                base_path = str(image_path).rsplit('.', 1)[0]
                fixed_path = f"{base_path}_fixed.jpg"
                
                self.logger.info(f"重新保存图片以确保兼容性: {fixed_path}")
                
                # 转换图片模式
                if mode == 'RGBA':
                    # 创建白色背景
                    rgb_img = Image.new('RGB', size, (255, 255, 255))
                    rgb_img.paste(img, mask=img.split()[-1])  # 使用alpha通道作为mask
                    img = rgb_img
                elif mode != 'RGB':
                    img = img.convert('RGB')
                
                # 重新保存为高质量JPEG
                img.save(fixed_path, 'JPEG', quality=95, optimize=True)
                
                # 验证新保存的图片
                with Image.open(fixed_path) as test_img:
                    test_img.verify()
                
                # 删除原文件
                if fixed_path != str(image_path):
                    os.remove(image_path)
                    self.logger.info(f"已删除原文件，保存为兼容格式: {fixed_path}")
                
                return fixed_path
                
        except Exception as e:
            self.logger.error(f"图片格式验证失败: {str(e)}")
            return None

    def format_publish_date(self, publish_time):
        """格式化发布日期为YYYY-MM-DD格式"""
        try:
            dt = None
            
            # 如果publish_time是时间戳（数字）
            if isinstance(publish_time, (int, float)):
                # 判断是秒级还是毫秒级时间戳
                if publish_time > 1e10:  # 毫秒级时间戳
                    dt = datetime.fromtimestamp(publish_time / 1000)
                else:  # 秒级时间戳
                    dt = datetime.fromtimestamp(publish_time)
                    
            # 如果publish_time是字符串，尝试解析
            elif isinstance(publish_time, str):
                # 先尝试作为时间戳解析
                try:
                    timestamp = float(publish_time)
                    if timestamp > 1e10:  # 毫秒级时间戳
                        dt = datetime.fromtimestamp(timestamp / 1000)
                    else:  # 秒级时间戳
                        dt = datetime.fromtimestamp(timestamp)
                except ValueError:
                    # 如果不是时间戳，尝试多种日期格式
                    date_formats = [
                        '%Y-%m-%d',
                        '%Y/%m/%d',
                        '%Y-%m-%d %H:%M:%S',
                        '%Y/%m/%d %H:%M:%S',
                        '%Y-%m-%dT%H:%M:%S',
                        '%Y-%m-%dT%H:%M:%SZ',
                        '%Y-%m-%dT%H:%M:%S.%fZ'
                    ]
                    
                    for fmt in date_formats:
                        try:
                            dt = datetime.strptime(publish_time, fmt)
                            break
                        except ValueError:
                            continue
                        
            # 如果publish_time是datetime对象
            elif isinstance(publish_time, datetime):
                dt = publish_time
                
            if dt is None:
                self.logger.warning(f"无法解析发布日期格式: {publish_time}")
                return None
            
            # 格式化日期为YYYY-MM-DD格式
            formatted_date = dt.strftime('%Y-%m-%d')
            
            return formatted_date
            
        except Exception as e:
            self.logger.error(f"格式化发布日期失败: {str(e)}")
            return None

    def generate_single_docx(self, article_info):
        """为单篇文章生成docx文件"""
        dd_article_id = article_info.get('dd_article_id')
        article_title = article_info.get('article_title', '未知标题')
        class_name = article_info.get('class_name', '未知课程')
        lecturer_name = article_info.get('lecturer_name', '未知讲师')
        
        # 查找对应的文章内容
        article_content_collection = self.db['Article_Content']
        content_doc = article_content_collection.find_one({'IdStr': dd_article_id})
        
        if not content_doc:
            self.logger.warning(f"未找到文章内容，ID: {dd_article_id}")
            return
        
        content = content_doc.get('content', '')
        
        # 创建文档
        doc = Document()
        
        # 设置标题
        title = doc.add_heading(article_title, level=1)
        title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        # 添加课程信息
        course_info = doc.add_paragraph()
        course_info.add_run(f"课程：{class_name}").bold = True
        course_info.add_run(f"\n讲师：{lecturer_name}").bold = True
        
        # 添加发布日期
        publish_time = article_info.get('publish_time')
        if publish_time:
            formatted_date = self.format_publish_date(publish_time)
            if formatted_date:
                course_info.add_run(f"\n发表日期：{formatted_date}").bold = True
        
        # 添加分隔线
        doc.add_paragraph("=" * 50)
        
        # 处理文章内容
        try:
            json_data = json.loads(content)
            texts = self.extract_text(json_data)
            self.add_content_to_document(doc, texts)
        except json.JSONDecodeError:
            self.logger.error(f"无法解析JSON内容，ID: {dd_article_id}")
            # 如果JSON解析失败，使用原始内容
            self.add_content_to_document(doc, content)
        
        # 创建输出目录
        output_dir = self.create_output_directory(class_name, lecturer_name)
        
        # 生成文件名
        safe_title = self.sanitize_filename(article_title)
        filename = f"{safe_title}_{dd_article_id}.docx"
        file_path = output_dir / filename
        
        # 保存文档
        doc.save(str(file_path))
        self.logger.info(f"已生成文档: {file_path}")
        
        # 标记为已处理
        self.mark_as_processed(article_info['_id'])

    def add_content_to_document(self, doc, content):
        """将内容添加到文档中"""
        if not content:
            doc.add_paragraph("暂无内容")
            return
        
        # 如果content是列表（从extract_text返回），直接处理
        if isinstance(content, list):
            for item in content:
                if item:
                    # 检查是否是标题（以【标题】开头）
                    if isinstance(item, str) and item.startswith("【标题】"):
                        title_text = item[3:]  # 移除"【标题】"前缀
                        doc.add_heading(title_text, level=2)
                    # 检查是否是图片
                    elif isinstance(item, dict) and item.get('type') == 'image':
                        success = self.add_image_to_document(doc, item['path'])
                        if not success:
                            self.logger.warning(f"图片添加失败，跳过: {item['path']}")
                    else:
                        # 检查是否包含结束文本
                        if isinstance(item, str) and self.should_terminate_generation(item):
                            self.logger.info("检测到结束文本，提前终止文档生成")
                            return
                        doc.add_paragraph(str(item))
            return
        
        # 如果content是字符串，按原来的方式处理
        if isinstance(content, str):
            # 简单的HTML标签处理
            clean_content = self.clean_html_content(content)
            
            # 按段落分割
            paragraphs = clean_content.split('\n')
            
            for para in paragraphs:
                para = para.strip()
                if para:
                    # 检查是否包含结束文本
                    if self.should_terminate_generation(para):
                        self.logger.info("检测到结束文本，提前终止文档生成")
                        return
                    
                    # 检查是否是标题（以数字开头或包含特定关键词）
                    if self.is_heading(para):
                        doc.add_heading(para, level=2)
                    else:
                        doc.add_paragraph(para)

    def clean_html_content(self, content):
        """清理HTML内容"""
        # 移除HTML标签
        content = re.sub(r'<[^>]+>', '', content)
        
        # 处理特殊字符
        content = content.replace('&nbsp;', ' ')
        content = content.replace('&amp;', '&')
        content = content.replace('&lt;', '<')
        content = content.replace('&gt;', '>')
        content = content.replace('&quot;', '"')
        
        # 清理多余的空行
        content = re.sub(r'\n\s*\n', '\n\n', content)
        
        return content.strip()

    def is_heading(self, text):
        """判断文本是否是标题"""
        # 以数字开头的段落可能是标题
        if re.match(r'^\d+[\.、]', text):
            return True
        
        # 包含特定关键词的可能是标题
        heading_keywords = ['章节', '第', '节', '部分', '主题', '要点']
        for keyword in heading_keywords:
            if keyword in text and len(text) < 50:
                return True
        
        return False

    def should_terminate_generation(self, text):
        """检测是否应该提前终止文档生成"""
        #termination_text = "今天广播就到这里。如果你觉得有帮助，欢迎分享转发给你的朋友。明天咱们接着聊AI"
        termination_text = "啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊"
        return termination_text in text

    def add_image_to_document(self, doc, image_path):
        """将图片添加到文档中"""
        try:
            self.logger.info(f"开始添加图片到文档: {image_path}")
            
            # 检查文件是否存在
            if not os.path.exists(image_path):
                self.logger.error(f"图片文件不存在: {image_path}")
                return False
            
            # 检查文件大小
            file_size = os.path.getsize(image_path)
            self.logger.info(f"图片文件大小: {file_size} 字节")
            
            if file_size == 0:
                self.logger.error(f"图片文件为空: {image_path}")
                return False
            
            # 检查文件是否为有效图片
            valid_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff']
            file_ext = os.path.splitext(image_path)[1].lower()
            if file_ext not in valid_extensions:
                self.logger.warning(f"图片文件扩展名可能不支持: {file_ext}, 路径: {image_path}")
            
            # 验证并修复图片格式
            validated_image_path = self.validate_and_fix_image(image_path)
            if not validated_image_path:
                self.logger.error(f"图片格式验证失败: {image_path}")
                return False
            
            # 添加图片到文档
            self.logger.info("创建段落和运行对象...")
            paragraph = doc.add_paragraph()
            run = paragraph.add_run()
            
            # 设置图片大小
            self.logger.info("添加图片到运行对象...")
            run.add_picture(validated_image_path, width=Inches(5))
            
            # 居中对齐
            self.logger.info("设置段落对齐...")
            paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            self.logger.info(f"图片已成功添加到文档: {image_path}")
            return True
                
        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            self.logger.error(f"添加图片到文档失败: {str(e)}")
            self.logger.error(f"错误详情: {error_details}")
            return False

    def create_output_directory(self, class_name, lecturer_name):
        """创建输出目录"""
        base_dir = Path(__file__).parent / 'output'
        base_dir.mkdir(exist_ok=True)
        
        # 创建子目录：课程名+讲师名
        safe_class_name = self.sanitize_filename(class_name)
        safe_lecturer_name = self.sanitize_filename(lecturer_name)
        sub_dir_name = f"{safe_class_name}_{safe_lecturer_name}"
        
        output_dir = base_dir / sub_dir_name
        output_dir.mkdir(exist_ok=True)
        
        return output_dir

    def sanitize_filename(self, filename):
        """清理文件名，移除非法字符"""
        # 移除或替换Windows文件名中的非法字符
        illegal_chars = r'[<>:"/\\|?*]'
        filename = re.sub(illegal_chars, '_', filename)
        
        # 限制长度
        if len(filename) > 100:
            filename = filename[:100]
        
        return filename.strip()

    def mark_as_processed(self, article_id):
        """标记文章为已处理"""
        try:
            article_info_collection = self.db['Article_Info']
            article_info_collection.update_one(
                {'_id': article_id},
                {'$set': {'processed': True, 'processed_time': datetime.now()}}
            )
            self.logger.info(f"已标记文章为已处理: {article_id}")
        except Exception as e:
            self.logger.error(f"标记文章失败: {str(e)}")

if __name__ == "__main__":
    generator = DocxGenerator()
    generator.generate_docx_files() 